#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
剪贴板图片OCR识别脚本
直接识别剪贴板中的截图，无需保存文件
"""

import win32clipboard
from PIL import Image
from io import BytesIO
import time
from OCR import wechat_ocr_from_memory, save_to_markdown
from colorama import init, Fore, Style

def get_clipboard_image():
    """从剪贴板获取图片"""
    try:
        # 打开剪贴板
        win32clipboard.OpenClipboard()
        
        # 检查剪贴板是否包含图片
        if win32clipboard.IsClipboardFormatAvailable(win32clipboard.CF_DIB):
            # 获取DIB格式的图片数据
            data = win32clipboard.GetClipboardData(win32clipboard.CF_DIB)
            
            # 将DIB数据转换为PIL Image
            # DIB数据需要添加文件头
            import struct
            
            # 创建BMP文件头
            bmp_header = struct.pack('<2sL4xL', b'BM', len(data) + 14, 14 + 40)
            bmp_data = bmp_header + data
            
            # 转换为PIL Image
            image = Image.open(BytesIO(bmp_data))
            
            win32clipboard.CloseClipboard()
            return image
            
        else:
            win32clipboard.CloseClipboard()
            return None
            
    except Exception as e:
        try:
            win32clipboard.CloseClipboard()
        except:
            pass
        print(f"获取剪贴板图片失败: {e}")
        return None

def get_clipboard_image_simple():
    """简化版剪贴板图片获取（使用PIL的ImageGrab）"""
    try:
        from PIL import ImageGrab
        
        # 直接从剪贴板获取图片
        image = ImageGrab.grabclipboard()
        
        if image is not None:
            return image
        else:
            print("剪贴板中没有图片")
            return None
            
    except Exception as e:
        print(f"获取剪贴板图片失败: {e}")
        return None

def ocr_clipboard_image():
    """识别剪贴板中的图片"""
    print(Fore.CYAN + "正在从剪贴板获取图片..." + Style.RESET_ALL)
    
    # 首先尝试简化版方法
    image = get_clipboard_image_simple()
    
    if image is None:
        # 如果简化版失败，尝试win32方法
        print("尝试使用win32方法获取剪贴板图片...")
        image = get_clipboard_image()
    
    if image is None:
        print(Fore.RED + "❌ 剪贴板中没有图片，请先截图并复制到剪贴板" + Style.RESET_ALL)
        return False
    
    print(Fore.GREEN + f"✅ 成功获取剪贴板图片，尺寸: {image.size}" + Style.RESET_ALL)
    
    # 进行OCR识别
    print(Fore.CYAN + "正在进行OCR识别..." + Style.RESET_ALL)
    
    try:
        ocr_results = wechat_ocr_from_memory(image)
        
        if not ocr_results:
            print(Fore.YELLOW + "⚠️ 未识别到任何文字" + Style.RESET_ALL)
            return False
        
        print(Fore.GREEN + f"✅ 识别成功！共识别到 {len(ocr_results)} 个文本块" + Style.RESET_ALL)
        
        # 显示识别结果
        print("\n" + "="*50)
        print(Fore.CYAN + "📝 识别结果:" + Style.RESET_ALL)
        print("="*50)
        
        for i, item in enumerate(ocr_results, 1):
            print(f"{Fore.YELLOW}文本块 {i}:{Style.RESET_ALL} {item['text']}")
            print(f"  坐标: 左:{item['left']:.1f}, 上:{item['top']:.1f}, 右:{item['right']:.1f}, 下:{item['bottom']:.1f}")
            print()
        
        # 保存结果
        timestamp = time.strftime("%Y%m%d_%H%M%S")
        output_file = f"clipboard_ocr_{timestamp}.md"
        save_to_markdown(ocr_results, output_file, f"剪贴板截图_{timestamp}")
        
        print(Fore.GREEN + f"✅ 识别结果已保存到: {output_file}" + Style.RESET_ALL)
        
        return True
        
    except Exception as e:
        print(Fore.RED + f"❌ OCR识别失败: {e}" + Style.RESET_ALL)
        return False

def main():
    """主函数"""
    init(autoreset=True)  # 初始化colorama
    
    print(Fore.CYAN + "🖼️  剪贴板图片OCR识别工具" + Style.RESET_ALL)
    print("="*50)
    
    while True:
        print(f"\n{Fore.YELLOW}请选择操作:{Style.RESET_ALL}")
        print("1. 识别剪贴板中的图片")
        print("2. 退出程序")
        
        choice = input(f"\n{Fore.CYAN}请输入选择 (1-2): {Style.RESET_ALL}").strip()
        
        if choice == '1':
            print("\n" + "-"*50)
            success = ocr_clipboard_image()
            if success:
                print(f"\n{Fore.GREEN}🎉 识别完成！{Style.RESET_ALL}")
            else:
                print(f"\n{Fore.RED}😞 识别失败{Style.RESET_ALL}")
            print("-"*50)
            
        elif choice == '2':
            print(f"\n{Fore.GREEN}👋 再见！{Style.RESET_ALL}")
            break
            
        else:
            print(f"\n{Fore.RED}❌ 无效选择，请重新输入{Style.RESET_ALL}")

if __name__ == '__main__':
    main() 